//=========================================================================
// Copyright (C) 2024 The C++ Component Model(COMO) Open Source Project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//=========================================================================

//#define _POSIX_C_SOURCE 200112L
#define _GNU_SOURCE
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <assert.h>
#include <sys/select.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <unistd.h>
#include <math.h>
#include <syslog.h>
#include <sched.h>
#include <linux/sched.h>
#include "zmtp.h"
#include "int64.h"
#include "math_utils.h"
#include "checksum.h"
#include "config.h"
#include "voter_message.h"
#include "voter_lib.h"
#include "voter_server.h"
#include "ZMQ_Switcher.h"
#include "rudp.h"

// global variables
// for rudp
Rudp rudp;
char msg_send[MAX_VoterMessage_SIZE];   // message buffer of send
char msg_recv[MAX_VoterMessage_SIZE];   // message buffer of recv
struct sockaddr_in serv_addr;           // server network info struct
struct sockaddr_in clie_addr;           // client network info struct
socklen_t clie_addr_size;

/**
 * Low level networking stuff
 * Here you will find basic socket stuff that should be part of
 * a decent standard C library, but you know... there are other
 * crazy goals for the future of C: like to make the whole language an
 * Undefined Behavior.
 */

/**
 * Set the specified socket in non-blocking mode, with no delay flag.
 */
int SocketSetNonBlockNoDelay(int fd)
{
    int flags;
    int yes = 1;

    /**
     * Set the socket nonblocking.
     * Note that fcntl(2) for F_GETFL and F_SETFL can't be
     * interrupted by a signal.
     */
    flags = fcntl(fd, F_GETFL);
    if (-1 == flags) {
        return -1;
    }
    if (fcntl(fd, F_SETFL, flags | O_NONBLOCK) == -1) {
        return -1;
    }

    // This is best-effort. No need to check for errors.
    setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &yes, sizeof(yes));
    return 0;
}

/**
 * Create a TCP server.
 * 
 * @param port The port number on which the server will listen.
 * @return returns fd of the created server;
 *         if an error occurs, returns -1.
 */
int CreateTCPServer(int port, const char *listenNs)
{
    int s;
    int reuse_flag = 1;
    int original_netns_fd;
    int target_netns_fd;
    struct sockaddr_in sa;

  do {
    if ((NULL != listenNs) && ('\0' != listenNs[0])) {
        // 保存当前network namespace的文件描述符
        original_netns_fd = open("/proc/self/ns/net", O_RDONLY);
        if (-1 == original_netns_fd) {
            perror("open /proc/self/ns/net");
            exit(EXIT_FAILURE);
        }

        // 打开目标network namespace
        char buf[PATH_MAX];
        snprintf(buf, sizeof(buf), "/var/run/netns/%s", listenNs);
        target_netns_fd = open(buf, O_RDONLY);
        if (-1 == target_netns_fd) {
            perror("open /var/run/netns/myns");
            close(original_netns_fd);
            exit(EXIT_FAILURE);
        }

        // 切换到目标namespace
        if (setns(target_netns_fd, CLONE_NEWNET) == -1) {
            perror("setns");
            close(original_netns_fd);
            close(target_netns_fd);
            exit(EXIT_FAILURE);
        }
    }

    // Create a TCP socket
    s = socket(AF_INET, SOCK_STREAM, 0);
    if (-1 == s) {
        break;
    }

    int enable  = 1;
    // 设置 SO_REUSEADDR 允许重用本地地址
    setsockopt(s, SOL_SOCKET, SO_REUSEADDR, &enable , sizeof(enable));
    // Linux 3.9+ 可以使用 SO_REUSEPORT 实现端口完全共享
    if (setsockopt(s, SOL_SOCKET, SO_REUSEPORT, &enable , sizeof(int)) < 0) {
        perror("setsockopt(SO_REUSEPORT) failed");
    }

    // Set socket option: address reuse (best-effort)
    setsockopt(s, SOL_SOCKET, SO_REUSEADDR, &reuse_flag, sizeof(reuse_flag));

    // Init the address structure
    memset(&sa, '\0', sizeof(sa));
    sa.sin_family = AF_INET;
    sa.sin_port = htons(port);
    sa.sin_addr.s_addr = htonl(INADDR_ANY);

    // Bind the socket to the specified port, then listen
    if ((bind(s, (struct sockaddr*)&sa, sizeof(sa)) == -1) ||
                                                        listen(s, 511) == -1) {
        close(s);
        s = -1;
        break;
    }
  } while (0);

    if((NULL != listenNs) && ('\0' != listenNs[0])) {
        // 切换回原始namespace
        if (setns(original_netns_fd, CLONE_NEWNET) == -1) {
            perror("setns back to original");
        }

        close(target_netns_fd);
        close(original_netns_fd);
    }

    return s;
}

/**
 * Create a UDP socket listening to 'port' ready to accept connections.
 */
int CreateUDPServer(int port, const char *listenNs)
{
    int s;
    int original_netns_fd;
    int target_netns_fd;

  do {
    if ((NULL != listenNs) && ('\0' != listenNs[0])) {
        // 保存当前network namespace的文件描述符
        original_netns_fd = open("/proc/self/ns/net", O_RDONLY);
        if (-1 == original_netns_fd) {
            perror("open /proc/self/ns/net");
            exit(EXIT_FAILURE);
        }

        // 打开目标network namespace
        char buf[PATH_MAX];
        snprintf(buf, sizeof(buf), "/var/run/netns/%s", listenNs);
        target_netns_fd = open(buf, O_RDONLY);
        if (-1 == target_netns_fd) {
            perror("open /var/run/netns/myns");
            close(original_netns_fd);
            exit(EXIT_FAILURE);
        }

        // 切换到目标namespace
        if (setns(target_netns_fd, CLONE_NEWNET) == -1) {
            perror("setns");
            close(original_netns_fd);
            close(target_netns_fd);
            exit(EXIT_FAILURE);
        }
    }

    // Create a UDP socket
    //s = socket(AF_INET, SOCK_DGRAM, 0);
    s = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (-1 == s) {
        break;
    }

    int enable  = 1;
    // 设置 SO_REUSEADDR 允许重用本地地址
    if (setsockopt(s, SOL_SOCKET, SO_REUSEADDR, &enable , sizeof(enable)) < 0) {
        perror("setsockopt(SO_REUSEADDR) failed");
    }

    enable = 1;
    // Linux 3.9+ 可以使用 SO_REUSEPORT 实现端口完全共享
    if (setsockopt(s, SOL_SOCKET, SO_REUSEPORT, &enable , sizeof(enable)) < 0) {
        perror("setsockopt(SO_REUSEPORT) failed");
    }

    enable = 1;
    // 开启广播权限
    if (setsockopt(s, SOL_SOCKET, SO_BROADCAST, &enable, sizeof(enable)) < 0) {
        perror("setsockopt(SO_BROADCAST) failed");
    }

    // Init the address structure
    memset(&serv_addr, '\0', sizeof(serv_addr));
    serv_addr.sin_family = AF_INET;
    serv_addr.sin_addr.s_addr = INADDR_ANY;
    serv_addr.sin_port = htons(port);

    // Bind the socket to the specified port
    if (bind(s, (const struct sockaddr*)&serv_addr, sizeof(serv_addr)) < 0) {
        close(s);
        s = -1;
        break;
    }

    // clie_addr and clie_addr_size are global variables. Is this necessary?
    // init clie_addr, clie_addr_size
    bzero(&clie_addr, sizeof(clie_addr));
    clie_addr_size = sizeof(struct sockaddr);

    // for rudp
    // init msg_recv & msg_send
    bzero(&msg_recv, sizeof(*msg_recv));
    bzero(&msg_send, sizeof(*msg_send));

  } while (0);

    if ((NULL != listenNs) && ('\0' != listenNs[0])) {
        // 切换回原始namespace
        if (setns(original_netns_fd, CLONE_NEWNET) == -1) {
            perror("setns back to original");
        }

        close(target_netns_fd);
        close(original_netns_fd);
    }

    return s;
}

int CreateUDPClient(int port)
{
    int s;
    //struct sockaddr_in server_addr;
    // now set server_addr as a global var
    // so comment out here

    // Create a UDP socket
    //s = socket(AF_INET, SOCK_DGRAM, 0);
    s = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (-1 == s) {
        return -1;
    }

    // Init the address structure
    // memset(&server_addr, 0, sizeof(server_addr));
    // server_addr.sin_family = AF_INET;
    // server_addr.sin_addr.s_addr = INADDR_ANY;
    // server_addr.sin_port = htons(port);
    // inet_pton(AF_INET, "127.0.0.1", &server_addr.sin_addr);
    memset(&serv_addr, '\0', sizeof(serv_addr));
    serv_addr.sin_family = AF_INET;
    serv_addr.sin_addr.s_addr = INADDR_ANY;
    serv_addr.sin_port = htons(port);

    // for rudp
    // init msg_recv & msg_send
    bzero(&msg_recv, sizeof(*msg_recv));
    bzero(&msg_send, sizeof(*msg_send));

    return s;
}

// func in rudp/comm.c
void udp_send(RudpPackage *p, int sock_fd, struct sockaddr_in serv_addr)
{
    // send package
    memcpy(msg_send, p->buffer, p->sz);
    size_t size_send = sendto(sock_fd, msg_send, p->sz, 0, 
                            (struct sockaddr *)(&serv_addr), sizeof(serv_addr));

    if (size_send <= 0) {
        if (-1 == close(sock_fd)) {
            syslog(LOG_ERR, "sendto() close() error");
        }
        syslog(LOG_ERR, "sendto() error");
    }
    printf("Sent package size = %zu\n", size_send);
}

void dump_recv(Rudp *U)
{
    char tmp[MAX_PACKAGE];
    int n;
    while ((n = RudpRecv(U, tmp))) {
        if (n < 0) {
            syslog(LOG_ERR, "RudpRecv, CORRPUT\n");
            break;
        }

        // Since codes below can only print VoterMessage (not Package) info, 
        // they are commented out.
        // printf("RECV ");
        // for (int i = 0; i < n; i++) {
        //     printf("%02x ", (uint8_t)tmp[i]);
        // }
        // printf("\n");
    }
}

/**
 * Create a TCP socket and connect it to the specified address.
 * On success the socket descriptor is returned, otherwise -1.
 *
 * If 'nonblock' is non-zero, the socket is put in nonblocking state
 * and the connect() attempt will not block as well, but the socket
 * may not be immediately ready for writing.
 */
int TCPConnect(const char *addr, int port, int nonblock)
{
    int s;
    int retval = -1;
    struct addrinfo hints, *servinfo, *p;

    char portstr[6];    // Max 16 bit number string length.
    snprintf(portstr, sizeof(portstr), "%d", port);
    memset(&hints, 0, sizeof(hints));
    hints.ai_family = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;

    if (getaddrinfo(addr, portstr, &hints, &servinfo) != 0) {
        return -1;
    }

    for (p = servinfo;  p != NULL;  p = p->ai_next) {
        /**
         * Try to create the socket and to connect it.
         * If we fail in the socket() call, or on connect(), we retry with
         * the next entry in servinfo.
         */
        s = socket(p->ai_family, p->ai_socktype, p->ai_protocol);
        if (-1 == s) {
            continue;
        }

        // Put in non blocking state if needed.
        if (nonblock && (SocketSetNonBlockNoDelay(s) == -1)) {
            close(s);
            break;
        }

        // Try to connect.
        if (connect(s, p->ai_addr, p->ai_addrlen) == -1) {
            /**
             * If the socket is non-blocking, it is ok for connect() to
             * return an EINPROGRESS error here.
             */
            if ((errno == EINPROGRESS) && nonblock) {
                return s;
            }

            // Otherwise it's an error.
            close(s);
            break;
        }

        /**
         * If we ended an iteration of the for loop without errors, we
         * have a connected socket. Let's return to the caller.
         */
        retval = s;
        break;
    }

    freeaddrinfo(servinfo);
    return retval;  // Will be -1 if no connection succeded.
}

// Create TCP connection
int createTCPConnection(const char *host, int port)
{
    int s = TCPConnect(host, port, 0);
    if (-1 == s) {
        perror("Connecting to server");
        exit(1);
    }
    return s;
}

/**
 * If the listening socket signaled there is a new connection ready to
 * be accepted, we accept(2) it and return -1 on error or the new client
 * socket on success.
 */
int AcceptClient(int server_socket, struct sockaddr_in *sa)
{
    int s;

    while (true) {
        socklen_t slen = sizeof(struct sockaddr_in);
        s = accept(server_socket, (struct sockaddr*)sa, &slen);
        if (-1 == s) {
            if (EINTR == errno) {
                continue;   //Try again.
            }
            else {
                return -1;
            }
        }
        break;
    }

    return s;
}

/**
 * Read N bytes from socket to buffer
 * This assumes buffer is at least bufferSize bytes long,
 * and that the socket is blocking.
 *
 * same function: s_tcp_recv (in zmtp_channel.c)
 */
int ReadNBytes(int socket, char* buffer, int bufferSize)
{
    int bytesRead = 0;
    int result;
    // Recursively read data until bufferSize bytes are read or an error occurs.
    while (bytesRead < bufferSize) {
        result = recv(socket, buffer + bytesRead, bufferSize - bytesRead, 0);
        /**
         * 在读取过程中，如果recv函数返回值小于等于0，表明socket连接可能已经断开，
         * 除非errno被设置为EINTR（这意味着recv函数是由于接收到信号而返回的，而不是因为连接断开）。
         * 此外，非阻塞socket可能返回EWOULDBLOCK或EAGAIN，表示当前没有数据可读，但连接并未断开。
         */
        // 发生错误且不是由于接收到信号或非阻塞读取导致的，返回-1
        // TODO: 由于未知原因，目前调试必定返回EWOULDBLOCK和EAGAIN，导致无法识别断开，故注释之
        // if ((result <= 0) && (errno != EINTR) 
        //                     && (errno != EWOULDBLOCK) && (errno != EAGAIN)) {
        if ((result <= 0) && (errno != EINTR)) {
            return -1;
        }
        // 累加实际读取的字节数
        bytesRead += result;
    }

    return bytesRead;
}

void printVoterMessage(VoterMessage *msg)
{
    printf("VoterMessage content:\n");
    printf("  Length         : %d\n", msg->len);
    printf("  CRC32          : %x\n", msg->crc32);
    printf("  Type           : %s\n", GetVoterMessageTypeName(msg->type));
    printf("  Reserved Chars : %s\n", msg->reservedChars);
    printf("  Transaction ID : %ld\n", msg->idTransaction);
    printf("  Lamport Clock  : %ld\n", msg->lamportClock);
    printf("  Data           : %ld\n", msg->data);
    printf("  Buffer         : %s\n", msg->buf);
    printf("---------------------------\n");
}

// A simplified version of printVoterMessage()
void printVoterMsg(const char *prompt, VoterMessage *msg)
{
    time_t timep;
    time(&timep); // 获取当前时间的秒数
    struct tm *p = gmtime(&timep); // 将秒数转换为 tm 结构体
    printf("%s [%02d:%02d:%02d %04d/%02d/%02d] ", prompt, 8 + p->tm_hour,
            p->tm_min, p->tm_sec, 1900 + p->tm_year, 1 + p->tm_mon, p->tm_mday);

    printf("Msg: len=%d, type=%s, id_trans=%ld, data=%ld, buf=%s, crc32=%x, TJ=%s\n",
           msg->len, GetVoterMessageTypeName(msg->type),
           msg->idTransaction, msg->data, 
           msg->buf, msg->crc32, msg->reservedChars);
}

/**
 * Send a VoterMessage through a specified socket using TCP or UDP.
 * 
 * @param s             Socket fd for sending the message.
 * @param type          The message type.
 * @param id_trans      Unique transaction ID.
 * @param lam_clock     Lamport clock value to ensure event ordering.
 * @param data          Numeric data to be sent with the message.
 * @param buf           Buffer containing additional message data.
 * @param rawDataSize   Size of the additional message data.
 * @param isTCP         Whether to use TCP (true) or UDP (false).
 * @return              Pointer to the VoterMessage structure that was sent.
 */
VoterMessage *sendVoterMessage(int s, char type, uint64_t id_trans, 
                              uint64_t lam_clock, int64_t data, char *buf,
                              int32_t rawDataSize, bool isTCP)
{
    // Construct VoterMessage for transmission
    char buffer[MAX_VoterMessage_SIZE];
    VoterMessage *box = BoxVoterMessage(buffer,
                MAX_VoterMessage_SIZE,  // bufSize
                type,                   // type
                id_trans,               // idTransaction
                lam_clock,              // lamportClock
                data,                   // data
                buf,
                rawDataSize);           // rawDataSize

    // Send message regarding the protocol (TCP/UDP)
    if (isTCP) {
        // Send the constructed VoterMessage to socket s
        int num = write(s, box, box->len);

#ifdef __DBEUG__
        printf("Written %d bytes:\n", num);
#endif
        // Check if bytes written are not as expected
        if (num != box->len) {
            printf("Error: (box->len == %d) bytes expected\n", box->len);
        }
        //print_hex((char*)box, box->len);
    }
    else { // UDP
        RudpSend(&rudp, (const char*)box, box->len);
        RudpPackage *p = RudpUpdate(&rudp, NULL, 0, 1);
        udp_send(p, s, serv_addr);

#ifdef __DBEUG__
        printf("Written %d bytes:\n", p->sz);
#endif
        //print_hex(p->buffer, p->sz);
    }
    // Show data package info
    //printVoterMessage(box);
#ifdef __DBEUG__
    printVoterMsg("sendVoterMessage", box);
#endif

    return box;
}

// Print bytes in the buffer (in hexadecimal form)
void print_hex(char *readbuf, int nread)
{
    for (int i = 0;  i < nread;  i++) {
        printf("%02X ", (unsigned char)readbuf[i]);
        if ((i % 8) == 7) {
            printf("\n");
        }
    }
    // Print \n in the last line if < 8 bytes
    if ((nread % 8) != 0) {
        printf("\n");
    }
}

/**
 * zmtp_channel.c _______________________________
 * // This is our greeting (64 octets)
 * const struct zmtp_greeting outgoing = {
 *     .signature = { 0xff, 0, 0, 0, 0, 0, 0, 0, 1, 0x7f },
 *     .version   = { 3, 0 },
 *     .mechanism = { 'N', 'U', 'L', 'L', '\0' }
 *  };
 */
static byte zmtp_signature[10] = { 0xffu, 0, 0, 0, 0, 0, 0, 0, 1, 0x7fu };

/**
 * WelcomeToServer
 */
int WelcomeToServer(VoterClient *c)
{
    if (c->fd < 0) {
        return -1;
    }
    printf("Client (fd = %d) connected.\n", c->fd);
    /**
     * Important: we don't do ANY BUFFERING. We just use the kernel
     * socket buffers. If the content does not fit, we don't care.
     * This is needed in order to keep this program simple.
     */
    size_t num = write(c->fd, zmtp_signature, sizeof(zmtp_signature));
    if (num != sizeof(zmtp_signature)) {
        return -2;
    }
    printf("Sent zmtp_signature to fd = %d.\n", c->fd);
    // print_hex((char*)zmtp_signature, sizeof(zmtp_signature));

    return 0;
}

/**
 * ResponseZMQ
 *
 * Called by clients after received Welcome message, 
 * to claim whether theirselves are ZMQ channels
 */
int ResponseZMQ(int fd, bool isZMQ)
{
    /**
     * zmtp_channel.c _______________________________
     * assert (incoming.signature[0] == 0xff);
     *
     * As long as the answer message is not 0xff, it is non-ZMQ
     */
    if (isZMQ) {
        char ch_zmq = 0xff;
        return write(fd, &ch_zmq, 1);
    }
    else {
        return write(fd, "T", 1);
    }
}

#define NANOSECONDS_PER_SECOND      1000000000
#define NANOSECONDS_PER_MILLISECOND 1000000
#define MILLISECONDS_PER_SECOND     1000

/**
 * NOTICE: ts refers to different meanings before/after it is set.
 * 
 * Judge whether delayTime is too short for OS to schedule (if so, return -1),
 * where delayTime == current time - start time (ts, before set)
 * If not, we set ts to (ts + nanoSeconds), and wait until the time point ts
 * 
 * @param ts           1. At first, it is start time
 *                     2. Later, it is set to (ts + nanoSeconds)
 * @param delayTime    Calculated delay time
 * @param nanoSeconds  Nanoseconds to wait
 * 
 * @return -1 if delayTime is too short, else return 0
 */
int TimespecWaitNanoseconds(struct timespec *ts, int64_t *delayTime,
                                                               long nanoSeconds)
{
    // Get current time
    struct timespec currentTime;
    clock_gettime(CLOCK_MONOTONIC, &currentTime);

    // If ts is NULL or the time < 0, use current time as start time
    if ((NULL == ts) || (ts->tv_sec <= 0)) {
        ts = &currentTime;
        *delayTime = 0;
    }
    else {
        long long int nsecTs = (long long int)(ts->tv_sec) *
                                        NANOSECONDS_PER_SECOND + ts->tv_nsec;
        long long int nsecCurrentTime = (long long int)(currentTime.tv_sec) *
                                NANOSECONDS_PER_SECOND + currentTime.tv_nsec;

        *delayTime = nsecCurrentTime - nsecTs;
        if (*delayTime < NANOSECONDS_FOR_SCHEDULE) {
            /**
             * The difference between ts and current time is too short to
             * schedule a clock_nanosleep(), so we don't sleep & return -1.
             */
            return -1;
        }
    }

    // set ts to (ts + nanoSeconds)
    long nsec = ts->tv_nsec + nanoSeconds;
    ts->tv_sec += nsec / NANOSECONDS_PER_SECOND;
    ts->tv_nsec = nsec % NANOSECONDS_PER_SECOND;

    // Even if interrupted by signal, continue waiting.
    // Explaination:
    // 1. If clock_nanosleep() returns 0, means we reached ts peacefully.
    // 2. If clock_nanosleep() returns -1 (interrupted by signal),
    //    since clock_nanosleep() is put in a while(), it will be called
    //    again, and try to wait until (the absolute time point) ts.
    
    while (clock_nanosleep(CLOCK_MONOTONIC, TIMER_ABSTIME, ts, NULL) == -1) {
        // If errno is not EINTR (indicates an error), stop waiting.
        if (errno != EINTR) {
            break;
        }
    }

    return 0;
}
